# Testing readout modules.
# 23-10-31

import os
import sys
import jax
import jax.numpy as jnp
import numpy as np
import haiku as hk
import mindspore as ms
import mindspore.nn as ms_nn

from jax import Array, jit
from typing import Optional, Union, List
from jax.nn.initializers import lecun_normal, normal, zeros, ones
from mindspore import load_param_into_net, Tensor, nn, context

sys.path.append(os.path.dirname(sys.path[0]))

# setting context & hyper parameters
context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")

from cybertron.readout import AtomwiseReadout, GraphReadout, PairwiseReadout
from ms_cybertron.readout import AtomwiseReadout as ms_AtomwiseReadout
from ms_cybertron.readout import GraphReadout as ms_GraphReadout
from ms_cybertron.readout import PairwiseReadout as ms_PairwiseReadout

EPSILON = 1e-3
rng = jax.random.PRNGKey(42)

# create and show data for test
node_vec = np.random.uniform(size=(3, 5, 4)) # (B, A, F)
node_mask = np.sum(node_vec, axis=-1) < 2 # (B, A)
edge_vec = np.random.uniform(size=(3, 5, 5, 4)) # (B, A, A, F)
edge_mask = np.sum(edge_vec, axis=-1) < 2 # (B, A, A)
edge_cutoff = np.random.uniform(size=(3, 5, 5)) # (B, A, A)

# test atomwise readout
print("=================Test AtomwiseReadout====================")
def atomwise_fn(x, _a, y, _b, _c, training=False):
    atomwise = AtomwiseReadout(dim_node_rep=4, activation='silu')(x, _a, y, _b, _c) # type: ignore
    return atomwise

atomwise_fn = jax.vmap(atomwise_fn, in_axes=(0, 0, 0, 0, 0))
ts_atomwise_fn = hk.transform(atomwise_fn, apply_rng=True)
atomwise_params = ts_atomwise_fn.init(rng, node_vec, edge_vec, node_mask, edge_mask, edge_cutoff)
atomwise_out = jit(ts_atomwise_fn.apply)(atomwise_params, rng, node_vec, edge_vec, node_mask, edge_mask, edge_cutoff)
print(f"atomwise_out: shape of {atomwise_out.shape}, value of \n{atomwise_out}")

_name = 'atomwise_readout/~/decoder/~/mlp/~/linear'
_params = hk.data_structures.to_mutable_dict(atomwise_params)
ms_atomwise_fn = ms_AtomwiseReadout(dim_node_rep=4, activation='silu',)
ms_params = {}
ms_params['decoder.output.mlp.0.weight'] = \
    ms.Parameter(Tensor(np.array(_params[f'{_name}_0']['w']).T, ms.float32))
ms_params['decoder.output.mlp.0.bias'] = \
    ms.Parameter(Tensor(np.array(_params[f'{_name}_0']['b']), ms.float32))
ms_params['decoder.output.mlp.1.weight'] = \
    ms.Parameter(Tensor(np.array(_params[f'{_name}_1']['w']).T, ms.float32))
ms_params['decoder.output.mlp.1.bias'] = \
    ms.Parameter(Tensor(np.array(_params[f'{_name}_1']['b']), ms.float32))
load_param_into_net(ms_atomwise_fn, ms_params)

ms_node_vec = Tensor(node_vec, ms.float32)
ms_node_mask = Tensor(node_mask, ms.bool_)
ms_atomwise_out = ms_atomwise_fn(node_rep=ms_node_vec, edge_rep=None, atom_mask=ms_node_mask)
print(f"ms_atomwise_out: shape of {ms_atomwise_out.shape}, value of \n{ms_atomwise_out}")

atomwise_same = np.allclose(atomwise_out, ms_atomwise_out.asnumpy(), atol=EPSILON)
print(f"atomwise readout is the same: {atomwise_same} in accuracy of {EPSILON}")
atomwise_diff = atomwise_out - ms_atomwise_out.asnumpy()
atomwise_norm = np.linalg.norm(atomwise_diff)
print(f"norm of diff: {atomwise_norm}")

# # test graph readout
# print("=================Test GraphReadout====================")
# def graph_fn(x, y, training=False):
#     graph = GraphReadout(dim_node_rep=4, activation='silu')(x, y) # type: ignore
#     return graph

# graph_fn = jax.vmap(graph_fn, in_axes=(0, 0,))
# ts_graph_fn = hk.transform(graph_fn, apply_rng=True)
# graph_params = ts_graph_fn.init(rng, node_vec, node_mask)
# graph_out = jit(ts_graph_fn.apply)(graph_params, rng, node_vec, node_mask)
# print(f"graph_out: shape of {graph_out.shape}, value of \n{graph_out}")

# _name = 'graph_readout/~/decoder/~/mlp/~/linear'
# _params = hk.data_structures.to_mutable_dict(graph_params)
# ms_graph_fn = ms_GraphReadout(dim_node_rep=4, activation='silu',)
# ms_params = {}
# ms_params['decoder.output.mlp.0.weight'] = \
#     ms.Parameter(Tensor(np.array(_params[f'{_name}_0']['w']).T, ms.float32))
# ms_params['decoder.output.mlp.0.bias'] = \
#     ms.Parameter(Tensor(np.array(_params[f'{_name}_0']['b']), ms.float32))
# ms_params['decoder.output.mlp.1.weight'] = \
#     ms.Parameter(Tensor(np.array(_params[f'{_name}_1']['w']).T, ms.float32))
# ms_params['decoder.output.mlp.1.bias'] = \
#     ms.Parameter(Tensor(np.array(_params[f'{_name}_1']['b']), ms.float32))
# load_param_into_net(ms_graph_fn, ms_params)

# ms_node_vec = Tensor(node_vec, ms.float32)
# ms_node_mask = Tensor(node_mask, ms.bool_)
# ms_graph_out = ms_graph_fn(node_rep=ms_node_vec, edge_rep=None, atom_mask=ms_node_mask)
# print(f"ms_graph_out: shape of {ms_graph_out.shape}, value of \n{ms_graph_out}")

# graph_same = np.allclose(graph_out, ms_graph_out.asnumpy(), atol=EPSILON)
# print(f"graph readout is the same: {graph_same} in accuracy of {EPSILON}")
# graph_diff = graph_out - ms_graph_out.asnumpy()
# graph_norm = np.linalg.norm(graph_diff)
# print(f"norm of diff: {graph_norm}")

# # test pairwise readout
# print("=================Test PairwiseReadout====================")
# def pairwise_fn(x, y, z, training=False):
#     pairwise = PairwiseReadout(dim_node_rep=4, dim_edge_rep=4, activation='silu')(x, y, z) # type: ignore
#     return pairwise

# pairwise_fn = jax.vmap(pairwise_fn, in_axes=(0, 0, 0,))
# ts_pairwise_fn = hk.transform(pairwise_fn, apply_rng=True)
# pairwise_params = ts_pairwise_fn.init(rng, edge_vec, edge_mask, edge_cutoff)
# pairwise_out = jit(ts_pairwise_fn.apply)(pairwise_params, rng, edge_vec, edge_mask, edge_cutoff)
# print(f"pairwise_out: shape of {pairwise_out.shape}, value of \n{pairwise_out}")

# _name = 'pairwise_readout/~/decoder/~/mlp/~/linear'
# _params = hk.data_structures.to_mutable_dict(pairwise_params)
# ms_pairwise_fn = ms_PairwiseReadout(dim_node_rep=4, dim_edge_rep=4, activation='silu',)
# ms_params = {}
# ms_params['decoder.output.mlp.0.weight'] = \
#     ms.Parameter(Tensor(np.array(_params[f'{_name}_0']['w']).T, ms.float32))
# ms_params['decoder.output.mlp.0.bias'] = \
#     ms.Parameter(Tensor(np.array(_params[f'{_name}_0']['b']), ms.float32))
# ms_params['decoder.output.mlp.1.weight'] = \
#     ms.Parameter(Tensor(np.array(_params[f'{_name}_1']['w']).T, ms.float32))
# ms_params['decoder.output.mlp.1.bias'] = \
#     ms.Parameter(Tensor(np.array(_params[f'{_name}_1']['b']), ms.float32))
# load_param_into_net(ms_pairwise_fn, ms_params)

# ms_edge_vec = Tensor(edge_vec, ms.float32)
# ms_edge_mask = Tensor(edge_mask, ms.bool_)
# ms_edge_cutoff = Tensor(edge_cutoff, ms.float32)
# ms_pairwise_out = ms_pairwise_fn(ms_edge_vec, ms_edge_mask, ms_edge_cutoff)
# print(f"ms_pairwise_out: shape of {ms_pairwise_out.shape}, value of \n{ms_pairwise_out}")

# pairwise_same = np.allclose(pairwise_out, ms_pairwise_out.asnumpy(), atol=EPSILON)
# print(f"pairwise readout is the same: {pairwise_same} in accuracy of {EPSILON}")
# pairwise_diff = pairwise_out - ms_pairwise_out.asnumpy()
# pairwise_norm = np.linalg.norm(pairwise_diff)
# print(f"norm of diff: {pairwise_norm}")


